# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/11 下午5:56
@Author yzpang
@Function: 
"""
from modelserver.train.paddle_train import TrainPaddle
from paddlenlp.transformers import AutoModel, AutoTokenizer


def train():
    model_path = r"E:\AI\framework\paddle\models\ernie-3.0-base-zh"
    # model_path = r"ernie-3.0-base-zh"
    label_path = r'E:\AI\datasets\paddle\hierarchical-class-train\label.txt'
    train_path = r'E:\AI\datasets\paddle\hierarchical-class-train\train.txt'
    dev_path = r'E:\AI\datasets\paddle\hierarchical-class-train\dev.txt'
    classification_type = "hierarchical-class"
    print('begin build trainer')
    trainer = TrainPaddle(model_path=model_path, classification_type=classification_type)
    print('build trainer success')
    trainer.train(label_path, train_path, dev_path)


def add_config():
    import time
    model_path = r"ernie-3.0-base-zh"
    save_path = r"E:\AI\framework\paddle\models\ernie-3.0-base-zh-new"
    model = AutoModel.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)

    time.sleep(10)


if __name__ == '__main__':
    train()
    # add_config()
